/*
 * Copyright (c) 2019, XXX All Rights Reserved.
 */
package net.sf.log4jdbc.agent;

import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;

import javassist.CtClass;
import net.sf.log4jdbc.agent.driver.JdbcDriverTransformer;
import net.sf.log4jdbc.agent.driver.JdbcPooledConnectionTransformer;


/**
 * <p>SQL 监控Agent</p>
 * @author zhanxb
 * @since 2021/8/30
 **/
public class SqlMonitorAgent
{
    public static void premain(String agentArg, Instrumentation instrumentation)
    {
        if (JdbcPooledConnectionTransformer.AGENT_TARGET.equalsIgnoreCase(agentArg))
        {
            instrumentation.addTransformer(new JdbcPooledConnectionTransformer());
        } else
        {
            instrumentation.addTransformer(new JdbcDriverTransformer());
        }

    }

    public static void agentmain(String agentArg, Instrumentation instrumentation)
    {
        try
        {
            Class[] loadedClasses = instrumentation.getAllLoadedClasses();
            for (Class loadedClass : loadedClasses)
            {
                String name = loadedClass.getName();
                if (JdbcDriverTransformer.isSupport(name))
                {
                    CtClass ctClass = JdbcDriverTransformer.enhanceConnectMethod(name);
                    if (ctClass!=null)
                    {
                        instrumentation.redefineClasses(new ClassDefinition(loadedClass, ctClass.toBytecode()));
                        System.out.println(name + "\u589e\u5f3a\u6210\u529f!");
                        return;
                    } else
                    {
                        System.out.println("\u672a\u8fdb\u884c\u4efb\u4f55\u589e\u5f3a!");
                    }
                }
            }
            System.out.println("\u672a\u8fdb\u884c\u4efb\u4f55\u589e\u5f3a!");
        } catch (Exception e)
        {
            e.printStackTrace();
        }
    }
}
